#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
多GPU服务器管理模块
支持动态添加/删除GPU服务器，负载均衡，健康检测
"""

import os
import json
import logging
import threading
import time
import requests
from typing import Dict, Any, List, Optional
from datetime import datetime

logger = logging.getLogger(__name__)

# 预设密码
WEBHOOK_SECRET = os.getenv('WEBHOOK_SECRET', 'gpu-server-register-to-api-gateway-2024')

class BackendServer:
    """GPU服务器信息"""
    def __init__(self, server_id: str, ip: str, port: int, 
                 internal_ip: str = None, public_ip: str = None, region: str = None):
        self.server_id = server_id
        self.ip = ip  # 保持向后兼容
        self.port = port
        self.url = f"http://{ip}:{port}"
        
        # 新增：支持内网IP和公网IP
        self.internal_ip = internal_ip or ip
        self.public_ip = public_ip
        self.region = region or 'unknown'
        self.server_url = f"http://{self.internal_ip}:{port}"  # 内网URL
        self.public_url = f"http://{self.public_ip}:{port}" if self.public_ip else None
        
        self.is_healthy = True
        self.fail_count = 0
        self.last_check_time = datetime.now()
        self.last_used_time = datetime.now()
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "server_id": self.server_id,
            "ip": self.ip,  # 保持向后兼容
            "port": self.port,
            "url": self.url,  # 保持向后兼容
            "internal_ip": self.internal_ip,
            "public_ip": self.public_ip,
            "region": self.region,
            "server_url": self.server_url,
            "public_url": self.public_url,
            "is_healthy": self.is_healthy,
            "fail_count": self.fail_count,
            "last_check_time": self.last_check_time.isoformat(),
            "last_used_time": self.last_used_time.isoformat()
        }

class BackendManager:
    """后端服务器管理器"""
    
    def __init__(self):
        self.servers: Dict[str, BackendServer] = {}
        self.current_index = 0
        self.health_check_interval = 3  # 3秒检测间隔
        self.max_fail_count = 5  # ✅ 改为5次失败后标记为不可用（更容错）
        self.auto_cleanup_threshold = 10  # ✅ 改为10次失败后自动删除（避免误删）
        self.health_check_thread = None
        self.running = False
        
        # 文件自动重新加载配置
        self.config_file = os.path.join(os.path.dirname(__file__), "backend_servers.json")
        self.last_file_mtime = 0  # 文件最后修改时间
        self.reload_check_interval = 1  # ✅ 优化为1秒检查（自动扩容响应快）
        self.last_reload_check = 0  # 上次检查时间
        self.lock = threading.RLock()  # 线程锁
        
        # 移除地域判断相关代码（不再需要）
        
        # 加载已保存的服务器列表
        self.load_servers()
        self.last_file_mtime = self._get_file_mtime()
        
        # 启动健康检测
        self.start_health_check()
    
    # 删除_get_c_server_region方法（不再需要地域判断）
    
    def _select_server_url(self, server: BackendServer) -> str:
        """选择服务器URL（强制使用公网IP）"""
        # C服务器是轻量应用服务器，G服务器是HAI实例
        # 它们使用不同的VPC，无法通过内网互通，必须使用公网IP
        if server.public_url:
            logger.debug(f"服务器{server.server_id}使用公网IP: {server.public_url}")
            return server.public_url
        else:
            # 如果没有公网IP，记录警告并回退到内网IP（可能无法连接）
            logger.warning(f"服务器{server.server_id}无公网IP，使用内网IP: {server.server_url}（可能无法连接）")
            return server.server_url
    
    def load_servers(self):
        """加载服务器列表"""
        try:
            config_file = os.path.join(os.path.dirname(__file__), "backend_servers.json")
            if os.path.exists(config_file):
                with open(config_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    for server_data in data:
                        server = BackendServer(
                            server_id=server_data['server_id'],
                            ip=server_data.get('ip', server_data.get('internal_ip')),
                            port=server_data['port'],
                            internal_ip=server_data.get('internal_ip'),
                            public_ip=server_data.get('public_ip'),
                            region=server_data.get('region', 'unknown')
                        )
                        self.servers[server.server_id] = server
                logger.info(f"加载了 {len(self.servers)} 台GPU服务器")
        except Exception as e:
            logger.error(f"加载服务器列表失败: {e}")
    
    def save_servers(self):
        """保存服务器列表"""
        try:
            config_file = os.path.join(os.path.dirname(__file__), "backend_servers.json")
            data = [server.to_dict() for server in self.servers.values()]
            with open(config_file, 'w', encoding='utf-8') as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
            logger.info("服务器列表已保存")
        except Exception as e:
            logger.error(f"保存服务器列表失败: {e}")
    
    def _get_file_mtime(self) -> float:
        """获取配置文件的修改时间"""
        try:
            if os.path.exists(self.config_file):
                return os.path.getmtime(self.config_file)
        except Exception:
            pass
        return 0
    
    def _check_and_reload_if_needed(self):
        """检查文件是否变更，如有变更则重新加载"""
        current_time = time.time()
        
        # 限制检查频率（避免每次请求都检查文件）
        if current_time - self.last_reload_check < self.reload_check_interval:
            return
        
        self.last_reload_check = current_time
        
        try:
            current_mtime = self._get_file_mtime()
            if current_mtime > self.last_file_mtime:
                logger.info(f"检测到backend_servers.json变更，重新加载（worker PID: {os.getpid()}）")
                with self.lock:
                    self.load_servers()
                    self.last_file_mtime = current_mtime
        except Exception as e:
            logger.error(f"检查文件变更失败: {e}")
    
    def verify_secret(self, secret: str) -> bool:
        """验证密码"""
        return secret == WEBHOOK_SECRET
    
    def generate_server_id(self) -> str:
        """自动生成服务器ID"""
        # 查找下一个可用的GPU编号
        existing_ids = set(self.servers.keys())
        for i in range(1, 1000):  # 最多支持999台服务器
            server_id = f"GPU-{i:03d}"
            if server_id not in existing_ids:
                return server_id
        raise Exception("服务器数量已达上限")
    
    def add_or_update_server(self, ip: str, port: int, secret: str, server_id: str = None,
                            internal_ip: str = None, public_ip: str = None, region: str = None) -> tuple[bool, str]:
        """添加或更新GPU服务器，返回(是否成功, server_id)"""
        try:
            # 验证密码
            if not self.verify_secret(secret):
                logger.warning(f"服务器密码验证失败")
                return False, ""
            
            # 使用internal_ip作为主要标识（如果提供）
            primary_ip = internal_ip or ip
            
            # 检查是否已存在相同IP的服务器
            existing_server = None
            for sid, server in self.servers.items():
                if (server.internal_ip == primary_ip and server.port == port) or \
                   (server.ip == primary_ip and server.port == port):
                    existing_server = sid
                    break
            
            if existing_server:
                # 更新现有服务器
                server = self.servers[existing_server]
                server.ip = primary_ip
                server.port = port
                server.internal_ip = internal_ip or primary_ip
                server.public_ip = public_ip
                server.region = region or 'unknown'
                server.url = f"http://{primary_ip}:{port}"
                server.server_url = f"http://{server.internal_ip}:{port}"
                server.public_url = f"http://{server.public_ip}:{port}" if server.public_ip else None
                server.is_healthy = True
                server.fail_count = 0
                logger.info(f"服务器 {existing_server} 已更新: 内网={server.internal_ip}, 公网={server.public_ip}, 地域={server.region}")
                server_id = existing_server
            else:
                # 添加新服务器，自动分配ID
                if not server_id:
                    server_id = self.generate_server_id()
                elif server_id in self.servers:
                    # 如果指定的ID已存在，自动生成新的
                    server_id = self.generate_server_id()
                
                server = BackendServer(
                    server_id=server_id,
                    ip=primary_ip,
                    port=port,
                    internal_ip=internal_ip or primary_ip,
                    public_ip=public_ip,
                    region=region or 'unknown'
                )
                self.servers[server_id] = server
                logger.info(f"服务器 {server_id} 已添加: 内网={server.internal_ip}, 公网={server.public_ip}, 地域={server.region}")
            
            # 保存配置
            self.save_servers()
            
            # 如果这是第一台服务器，启动健康检测
            if len(self.servers) == 1 and not self.running:
                self.start_health_check()
            
            return True, server_id
        except Exception as e:
            logger.error(f"添加/更新服务器失败: {e}")
            return False, ""
    
    def remove_server(self, server_id: str) -> bool:
        """删除GPU服务器"""
        try:
            if server_id in self.servers:
                del self.servers[server_id]
                logger.info(f"服务器 {server_id} 已删除")
                self.save_servers()
                
                # 如果没有服务器了，停止健康检测
                if len(self.servers) == 0:
                    self.stop_health_check()
                
                return True
            return False
        except Exception as e:
            logger.error(f"删除服务器失败: {e}")
            return False
    
    def get_next_server(self) -> Optional[BackendServer]:
        """获取下一个可用的服务器（负载均衡），自动选择合适的URL"""
        # 检查文件变更并自动重新加载
        self._check_and_reload_if_needed()
        
        if not self.servers:
            return None
        
        # 获取所有健康的服务器
        healthy_servers = [s for s in self.servers.values() if s.is_healthy]
        
        if not healthy_servers:
            logger.warning("没有可用的GPU服务器")
            return None
        
        # 顺序选择服务器（轮询）
        server = healthy_servers[self.current_index % len(healthy_servers)]
        self.current_index += 1
        server.last_used_time = datetime.now()
        
        # ✅ 添加详细日志
        logger.info(f"[负载均衡] 选择服务器: {server.server_id} ({server.ip}:{server.port}), 健康服务器总数: {len(healthy_servers)}, 当前索引: {self.current_index-1}")
        
        # 智能选择URL（内网优先）并更新server.url
        selected_url = self._select_server_url(server)
        server.url = selected_url
        
        return server
    
    def check_server_health(self, server: BackendServer) -> bool:
        """检查单个服务器健康状态"""
        try:
            # 健康检查也使用公网IP（与get_next_server保持一致）
            health_check_url = self._select_server_url(server)
            health_url = f"{health_check_url}/health"
            
            # 如果是HAI优化系统，使用不同的健康检查路径
            if server.server_id.startswith('HAI-OPT') or '127.0.0.1:8080' in health_check_url:
                health_url = f"{health_check_url}/api/v1/health"
            
            # 尝试访问健康检查端点
            response = requests.get(health_url, timeout=10)  # ✅ 改为10秒超时（更容错）
            
            if response.status_code == 200:
                # 验证响应内容
                try:
                    health_data = response.json()
                    
                    # 检查GPU是否可用（可选，严格模式）
                    if 'gpu_available' in health_data and not health_data.get('gpu_available'):
                        logger.warning(f"服务器 {server.server_id}: GPU不可用")
                        server.fail_count += 1
                        server.is_healthy = False
                        server.last_check_time = datetime.now()
                        
                        # 自动清理僵尸服务器
                        if server.fail_count >= self.auto_cleanup_threshold:
                            logger.warning(f"服务器{server.server_id}连续失败{server.fail_count}次，自动注销")
                            with self.lock:
                                self.servers.pop(server.server_id, None)
                                self.save_servers()
                        
                        return False
                    
                    logger.debug(f"服务器 {server.server_id} 健康检查通过: {health_data}")
                except Exception as json_err:
                    logger.debug(f"健康检查响应不是JSON格式: {json_err}，假定健康")
                
                # 健康检查成功
                server.fail_count = 0
                server.is_healthy = True
                server.last_check_time = datetime.now()
                
                if not server.is_healthy:
                    logger.info(f"服务器 {server.server_id} 重新上线")
                
                return True
            else:
                server.fail_count += 1
                logger.warning(f"服务器 {server.server_id} 健康检查失败: HTTP {response.status_code}")
                server.is_healthy = False
                server.last_check_time = datetime.now()
                
                # 自动清理僵尸服务器
                if server.fail_count >= self.auto_cleanup_threshold:
                    logger.warning(f"服务器{server.server_id}连续失败{server.fail_count}次，自动注销")
                    with self.lock:
                        self.servers.pop(server.server_id, None)
                        self.save_servers()
                
                return False
                
        except Exception as e:
            server.fail_count += 1
            server.is_healthy = False
            server.last_check_time = datetime.now()
            logger.warning(f"服务器 {server.server_id} 健康检查异常: {e}")
            
            # 自动清理僵尸服务器
            if server.fail_count >= self.auto_cleanup_threshold:
                logger.warning(f"服务器{server.server_id}连续失败{server.fail_count}次，自动注销")
                with self.lock:
                    self.servers.pop(server.server_id, None)
                    self.save_servers()
            
            return False
    
    def health_check_loop(self):
        """健康检查循环"""
        while self.running:
            try:
                if self.servers:
                    for server in list(self.servers.values()):
                        self.check_server_health(server)
                
                time.sleep(self.health_check_interval)
            except Exception as e:
                logger.error(f"健康检查循环出错: {e}")
                time.sleep(self.health_check_interval)
    
    def start_health_check(self):
        """启动健康检查"""
        if self.servers and not self.running:
            self.running = True
            self.health_check_thread = threading.Thread(target=self.health_check_loop, daemon=True)
            self.health_check_thread.start()
            logger.info("健康检查已启动")
    
    def stop_health_check(self):
        """停止健康检查"""
        self.running = False
        if self.health_check_thread:
            self.health_check_thread.join(timeout=5)
        logger.info("健康检查已停止")
    
    def mark_server_failed(self, server_id: str):
        """立即标记服务器为失败（用于请求失败后的快速响应）"""
        if server_id in self.servers:
            server = self.servers[server_id]
            server.fail_count += 1
            server.is_healthy = False
            server.last_check_time = datetime.now()
            logger.warning(f"服务器 {server_id} 标记为失败 (fail_count: {server.fail_count})")
            
            # 如果失败次数过多，自动删除
            if server.fail_count >= self.auto_cleanup_threshold:
                logger.warning(f"服务器 {server_id} 连续失败 {server.fail_count} 次，自动注销")
                with self.lock:
                    self.servers.pop(server_id, None)
                    self.save_servers()
    
    def get_next_server_with_retry(self, max_retries: int = None) -> Optional[BackendServer]:
        """
        获取下一个可用服务器（带智能重试）
        
        这个方法封装了故障转移逻辑：
        - 首次调用 get_next_server() 获取一个服务器
        - 如果业务层报告失败，会自动尝试下一个
        - 最多尝试所有健康服务器
        
        参数:
            max_retries: 最大重试次数，None表示尝试所有健康服务器
        
        返回:
            BackendServer对象，如果所有服务器都不可用则返回None
        """
        if max_retries is None:
            # 默认尝试所有健康服务器
            max_retries = len([s for s in self.servers.values() if s.is_healthy])
        
        if max_retries == 0:
            logger.error("[负载均衡] 没有可用的服务器")
            return None
        
        logger.info(f"[负载均衡] 开始选择服务器，最多尝试 {max_retries} 次")
        
        for attempt in range(1, max_retries + 1):
            server = self.get_next_server()
            if server:
                logger.info(f"[负载均衡] 第 {attempt} 次尝试：选择 {server.server_id} ({server.ip}:{server.port})")
                return server
            else:
                logger.warning(f"[负载均衡] 第 {attempt} 次尝试失败：没有可用服务器")
        
        logger.error(f"[负载均衡] 尝试了 {max_retries} 次，所有服务器都不可用")
        return None

    def get_all_servers(self) -> List[Dict[str, Any]]:
        """获取所有服务器信息"""
        return [server.to_dict() for server in self.servers.values()]
    
    def get_server_urls(self) -> List[str]:
        """获取所有健康服务器的URL列表"""
        return [server.url for server in self.servers.values() if server.is_healthy]
    
    def get_stats(self) -> Dict[str, Any]:
        """获取统计信息"""
        total = len(self.servers)
        healthy = len([s for s in self.servers.values() if s.is_healthy])
        unhealthy = total - healthy
        
        return {
            "total_servers": total,
            "healthy_servers": healthy,
            "unhealthy_servers": unhealthy,
            "health_check_running": self.running,
            "health_check_interval": self.health_check_interval,
            "max_fail_count": self.max_fail_count,
            "auto_cleanup_threshold": self.auto_cleanup_threshold,
            "servers": self.get_all_servers()
        }

# 全局实例
backend_manager = BackendManager()

